"""
数据预处理与加载
功能：将CSV数据转换为PyTorch Dataset
特点：支持多模态输入（波形+结构参数）
"""
import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import OneHotEncoder, MinMaxScaler
import json
import joblib
import os


class ComponentDataset(Dataset):
    def __init__(self, csv_path, window_size=1000):
        """
        参数说明：
        - csv_path: 数据集路径
        - window_size: 输入波形长度（默认1000点/10秒）
        """
        df = pd.read_csv(csv_path)

        # 结构参数预处理
        self._preprocess_structural_params(df)

        # 波形数据处理
        self.waves = self._process_waveforms(df)

        # 标签
        self.labels = df['damage_index'].values.astype(np.float32)

        dataset = ComponentDataset('data/component_data.csv')
        dataset.save_preprocessors('outputs')  # 此行必须执行

    def _preprocess_structural_params(self, df):
        """结构参数编码处理"""
        # 数值型参数归一化
        num_features = ['beam_width', 'beam_height', 'rebar_ratio']
        self.scaler = MinMaxScaler()
        df[num_features] = self.scaler.fit_transform(df[num_features])

        # 类别型参数独热编码
        self.encoder = OneHotEncoder(sparse_output=False)
        encoded_grade = self.encoder.fit_transform(df[['concrete_grade']])
        # 合并特征
        self.struct_params = np.hstack([
            df[num_features].values,
            encoded_grade
        ]).astype(np.float32)

    def _process_waveforms(self, df):
        """处理波形数据"""
        # 将JSON字符串解析为Python列表
        df['acc_x'] = df['acc_x'].apply(json.loads)
        df['acc_y'] = df['acc_y'].apply(json.loads)
        df['acc_z'] = df['acc_z'].apply(json.loads)

        # 转换为numpy数组（形状：样本数 × 3通道 × 1000点）
        waves = np.stack([
            np.array(df['acc_x'].tolist()),
            np.array(df['acc_y'].tolist()),
            np.array(df['acc_z'].tolist())
        ], axis=1)

        # 归一化到[-1,1]
        max_vals = np.max(np.abs(waves), axis=2, keepdims=True)
        return (waves / max_vals).astype(np.float32)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return {
            'wave': torch.tensor(self.waves[idx]),
            'struct': torch.tensor(self.struct_params[idx]),
            'label': torch.tensor(self.labels[idx])
        }

    def save_preprocessors(self, save_dir):
        """保存 scaler 和 encoder"""
        os.makedirs(save_dir, exist_ok=True)
        joblib.dump(self.scaler, os.path.join(save_dir, 'scaler.pkl'))
        joblib.dump(self.encoder, os.path.join(save_dir, 'encoder.pkl'))


def get_dataloader(csv_path, batch_size=32, shuffle=True):
    """创建数据加载器"""
    dataset = ComponentDataset(csv_path)
    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)